;; ===========
;; TERMINOLOGY
;; ===========

;; DIE: six sided die

;; DICE: more than one die

;; ROUND: rolling dice for attackers and defenders once

;; FIGHT: rolling dice for as many times as it takes for the attacker or
;; defender to be unable to roll dice

;; BATTLE: multiple subsequent fights, each with remaining attackers and
;; defenders after the previous fight

;; ANNIHILATION: rolling dice, playing rounds, until the attacker cannot roll
;; any die any longer or won the fight.

(use-modules
 ;; SRFI 1 for additional list procedures
 (srfi srfi-1)
 ;; SRFI 8 for receive form
 (srfi srfi-8)
 ;; SRFI-27 for random number utilities
 (srfi srfi-27)
 ;; SRFI-69 for hash tables
 (srfi srfi-69)
 ;; format for formatting floats
 (ice-9 format)
 ;; for functional structs (not part of srfi-9 directly)
 (srfi srfi-9 gnu)
 (rnrs enums)

 (display-utils)
 (helpers)
 (model)
 (random))


;; ==========
;; RISK LOGIC
;; ==========
(define roll-die
  (lambda* (rand-int-proc #:key (die-sides 6))
    (+ (rand-int-proc die-sides) 1)))


(define roll-dice
  (λ (dice-count rand-int-proc)
    (let loop ([remaining-dice dice-count]
               [results '()])
      (cond
       [(= remaining-dice 0) results]
       [else
        (loop (- remaining-dice 1)
              (cons (roll-die rand-int-proc) results))]))))


(define calculate-round-result
  (lambda* (fight-situation
            rolls-attacker
            rolls-defender
            #:key
            (rules *default-risk-rules*))
    ;; The defender has at maximum the same number of dice
    ;; the attack has, so we can iterate over the defender
    ;; rolls.
    (let loop ([remaining-rolls-att (sort rolls-attacker >)]
               [remaining-rolls-def (sort rolls-defender >)]
               [att-wins 0]
               [def-wins 0])
      (cond
       ;; Attacker rolls can be exhausted before defender
       ;; rolls, as it is possible to do a minority attack.
       [(or (null? remaining-rolls-def) (null? remaining-rolls-att))
        (make-fight-situation (+ (get-att-wins fight-situation) att-wins)
                              (- (get-att-rems fight-situation) def-wins)
                              (+ (get-def-wins fight-situation) def-wins)
                              (- (get-def-rems fight-situation) att-wins))]
       [else
        (let ([max-att-roll (car remaining-rolls-att)]
              [max-def-roll (car remaining-rolls-def)])
          (cond
           [(< max-att-roll max-def-roll)
            (loop (cdr remaining-rolls-att)
                  (cdr remaining-rolls-def)
                  att-wins
                  (+ def-wins 1))]
           [(> max-att-roll max-def-roll)
            (loop (cdr remaining-rolls-att)
                  (cdr remaining-rolls-def)
                  (+ att-wins 1)
                  def-wins)]
           ;; The case of equal eyes on the dice needs to be
           ;; treated specially, according to the rules
           ;; give.
           [(= max-att-roll max-def-roll)
            (let ([eq-eyes-win-rule (get-eq-eyes-win rules)])
              (cond
               [(eq? eq-eyes-win-rule 'def)
                (loop (cdr remaining-rolls-att)
                      (cdr remaining-rolls-def)
                      att-wins
                      (+ def-wins 1))]
               [(eq? eq-eyes-win-rule 'att)
                (loop (cdr remaining-rolls-att)
                      (cdr remaining-rolls-def)
                      (+ att-wins 1)
                      def-wins)]
               [(eq? eq-eyes-win-rule 'both)
                (loop (cdr remaining-rolls-att)
                      (cdr remaining-rolls-def)
                      (+ att-wins 1)
                      (+ def-wins 1))]
               [(eq? eq-eyes-win-rule 'none)
                (loop (cdr remaining-rolls-att)
                      (cdr remaining-rolls-def)
                      att-wins
                      def-wins)]
               [else
                (error "unknown equal eyes win rule")]))]))]))))


(define simulate-round
  (lambda* (fight-situation
            rand-int-proc
            #:key
            (rules *default-risk-rules*))
    "Roll dice and calculate the result of the round. Return a new fight
situation."

    (define calc-max-att-count
      (λ (att-dice-count rules)
        (min
         ;; At least one unit needs to remain in the source
         ;; area.
         (- att-dice-count 1)
         ;; Only so many may attack in one round, accordung
         ;; to the rules.
         (get-att-dice rules))))

    (define calc-max-def-count
      (λ (def-dice-count rules)
        (min
         ;; The defender may use up all their units.
         def-dice-count
         ;; Only so many may defend in one round, according
         ;; to the rules.
         (get-def-dice rules))))

    (let ([rolls-attacker
           (roll-dice (calc-max-att-count (get-att-rems fight-situation) rules)
                      rand-int-proc)]
          [rolls-defender
           (roll-dice (calc-max-def-count (get-def-rems fight-situation) rules)
                      rand-int-proc)])
      (calculate-round-result fight-situation
                              rolls-attacker
                              rolls-defender
                              #:rules rules))))


(define calc-att-winning-prob
  (λ (count-table)
    "Calculate the winning probability of the attacker given
a fight results table."
    (let ([att-wins-def-wins
           (hash-table-fold count-table
                            (λ (fight-result-as-list count previous)
                              ;; There are only 2 possible
                              ;; outcomes of
                              ;; annihilation. Either the
                              ;; remaining attackers are 1
                              ;; (NOTE/TODO: perhaps subject
                              ;; to rules), or the remaining
                              ;; defenders are 0. Any other
                              ;; result is not a result of
                              ;; annihilation.
                              (cond
                               ;; If the defender lost ...
                               [(= (fight-sit-list-def-rems fight-result-as-list) 0)
                                (cons (+ (car previous) count)
                                      (cdr previous))]
                               ;; If the attacker lost ...
                               [(> (fight-sit-list-def-rems fight-result-as-list) 0)
                                (cons (car previous)
                                      (+ (cdr previous) count))]
                               [else
                                (error
                                 "calc-att-winning-prob given non-annihilation fight results count table")]))
                            (cons 0 0))])
      (/ (car att-wins-def-wins)
         (+ (car att-wins-def-wins)
            (cdr att-wins-def-wins))))))


(define calc-def-win-prob
  (λ (count-table)
    "Calculate the winning probability of the defender, given an annihilation
fight result table."
    (- 1 (calc-att-winning-prob count-table))))


(define calc-consecutive-att-win-prob
  (lambda* (att-counts
            def-counts
            n-times
            rand-int-proc
            #:key
            (rules *default-risk-rules*))
    "Calculate the probability for consecutively winning mutliple fights of the
given counts of attackers against the given counts of defenders. Expects the
number of counts of attackers to be equal to the number of counts of
defenders (pairings)."
    (let loop ([rem-att-counts att-counts]
               [rem-def-counts def-counts]
               [prob 1])
      (cond
       [(null? rem-def-counts) prob]
       [else
        (loop
         (cdr rem-att-counts)
         (cdr rem-def-counts)
         (* prob
            (calc-att-winning-prob
             (try-n-times (make-init-fight-situation (car rem-att-counts)
                                                     (car rem-def-counts))
                          n-times
                          simulate-fight
                          rand-int-proc
                          #:rules rules))))]))))


(define try-n-times-general
  (lambda (try-proc n-times)
    "Try a try-proc n times and give back counted fight
results. try-proc takes as arguments a fight situation and
the random generator and the rules. Return a hash-table
using fight results as keys and count of those fight results
occurring as values."
    (let ([count-table (make-hash-table)])
      (let loop ([times-tried 0])
        (cond
         [(< times-tried n-times)
          (let* ([round-result (try-proc)]
                 [hash-key (fight-situation->hash-key round-result)])
            (hash-table-set!
             count-table
             hash-key
             (+ (hash-table-ref count-table hash-key (λ () 0)) 1))
            (loop (+ times-tried 1)))]
         [else
          count-table])))))


(define simulate-battle
  (lambda* (att-count
            def-counts
            rand-int-proc
            #:key
            (rules *default-risk-rules*)
            (verbose #f))
    "Simulate a battle exactly once given the attacker count and a list of
defender counts."
    (let iter-fights ([available-att-count att-count]
                      [def-counts def-counts]
                      [battle-result
                       (make-init-fight-situation att-count (apply + def-counts))]
                      [cur-fight-num 0])
      (when verbose
        (debug "fight number:" cur-fight-num))
      (when verbose
        (debug "rem. att:" available-att-count))
      (cond
       ;; If there are no defenders left to defend or no attackers left to
       ;; attack, then the battle is over and we need to return a fight
       ;; situation (result).
       [(null? def-counts)
        (when verbose
          (debug "STOP: no more defender groups"))
        battle-result]
       [(< available-att-count 2)
        (when verbose
          (debug "STOP: insufficient attackers:" available-att-count))
        battle-result]
       ;; Otherwise ...
       [else
        (when verbose
          (debug "rem. def:" (car def-counts)))
        (when verbose
          (debug "battle result so far:" battle-result))
        (let* ([fight-result
                (simulate-fight
                 (make-init-fight-situation available-att-count (car def-counts))
                 rand-int-proc
                 #:rules rules)]
               [updated-battle-result
                (make-fight-situation
                 (+ (get-att-wins battle-result)
                    (get-att-wins fight-result))
                 ;; NOTE: FUTURE: Future rules may contain a requirement of leaving
                 ;; more than 1 unit in the source area, which would require a change
                 ;; of this code.
                 (get-att-rems fight-result)
                 (+ (get-def-wins battle-result)
                    (get-def-wins fight-result))
                 (+ (get-def-rems fight-result)
                    (apply + (cdr def-counts))))])
          (when verbose
            (debug "fight gave following result:" fight-result))
          (when verbose
            (debug "updated battle result:" updated-battle-result))
          ;; Continue with the next battle, if there are defenders or attackers
          ;; left to fight.
          ;; NOTE: FUTURE: Future rules may contain a requirement of leaving
          ;; more than 1 unit in the source area, which would require a change
          ;; of this code.
          (iter-fights (- (get-att-rems updated-battle-result) 1)
                       (cdr def-counts)
                       updated-battle-result
                       (+ cur-fight-num 1)))]))))


(define calc-battle-att-win-prob
  (lambda* (att-count
            def-counts
            n-times
            rand-int-proc
            #:key
            (rules *default-risk-rules*)
            (verbose #f))
    "Calculate the probability for consecutively winning mutliple fights of the
given count of attackers against the given counts of defenders."
    ;; TODO: remove after removing debug of depth
    (define max-depths (length def-counts))

    ;; (1) Number of attackers changes at each fight, by at least 1.
    ;; (2) At each fight, there is another list of remaining defender
    ;; counts.
    ;; (3) The initial previous probability is 1, because it is the neutral
    ;; element of multiplication.
    (define iter
      (λ (att-count def-counts prev-prob branch)
        ;; TOOD: remove debug
        (when verbose
          (debug "now at branch" branch "with branch selection probability" (format-fraction prev-prob #:precision 5)))
        (cond
         ;; If all defenders are gone, return simply the
         ;; previous probability, which is the probability
         ;; for reaching this fight result from the fight
         ;; result before the previous fight situation.
         [(null? def-counts)
          ;; TOOD: remove debug
          (when verbose
            (debug "no defenders left, bubbling up to branch" (take branch (- (length branch) 1))))
          prev-prob]
         [else
          ;; TOOD: remove debug
          (when (= (- max-depths (length def-counts)) 1)
            (when verbose (debug "----------------------------------------")))

          (let ([count-table
                 ;; Calculate fight result probabilities for the current fight
                 ;; situation.
                 (try-n-times (make-init-fight-situation att-count (car def-counts))
                              n-times
                              simulate-fight
                              rand-int-proc
                              #:rules rules)])
            ;; Along the branches of a probability tree we multiply. Or in other
            ;; words probabilities for events associated with AND need to be
            ;; multiplied. The probabilities of battles after the previous fight
            ;; depend on that previous fight happening, so they are AND
            ;; associated.
            (* prev-prob
               ;; Look at each result using a fold over the entries of the hash
               ;; table of fight results.
               (hash-table-filter-fold count-table
                                       ;; Filter for fights won by the attacker.
                                       (λ (key val)
                                         (fight-res-list-att-win? key))
                                       ;; Sum all battles from won fight
                                       ;; results. Use addition as accumulation
                                       ;; procedure.
                                       ;; TOOD: remove debug logs, just use + instead of lambda
                                       (lambda (a b)
                                         (when verbose
                                           (debug "in branch" branch
                                                  "accumulating"
                                                  (format-fraction a #:precision 7) "+" (format-fraction b #:precision 7) "=" (format-fraction (+ a b) #:precision 7)))
                                         (+ a b))
                                       ;; For each fight result, which
                                       ;; represents a win for the attacker, we
                                       ;; need to calculate the battles after
                                       ;; it.
                                       (λ (key val)
                                         (iter
                                          ;; The attacker can only attack with
                                          ;; the remaining units. 1 attacker
                                          ;; must remain in the conquered area.
                                          ;; NOTE/TODO: If the rules allow empty
                                          ;; source areas, this part needs a
                                          ;; change.
                                          (- (fight-sit-list-att-rems key) 1)
                                          ;; The attacker annihilated one group
                                          ;; of defenders, so only give the
                                          ;; other defender counts for the next
                                          ;; battle calculation.
                                          (cdr def-counts)
                                          ;; Probability of the current fight
                                          ;; result, for which the subsequent
                                          ;; probability of winning the battle
                                          ;; is calculated.
                                          (/ val n-times)
                                          (append branch (list key))))
                                       ;; Add probabilities to 0, as 0 is the
                                       ;; neutral element of addition.
                                       0)))])))
    (iter att-count def-counts 1 '())))


(define try-n-times
  (lambda* (fight-situation
            n
            try-proc
            rand-int-proc
            #:key
            (rules *default-risk-rules*))
    "Try a try-proc n times and give back counted fight
results. try-proc takes as arguments a fight situation and
the random generator and the rules. Return a hash-table
using fight results as keys and count of those fight results
occurring as values."
    (let ([count-table (make-hash-table)])
      (let loop ([times-tried 0])
        (cond
         [(< times-tried n)
          (let* ([round-result
                  (try-proc fight-situation rand-int-proc #:rules rules)]
                 [hash-key (fight-situation->hash-key round-result)])
            (hash-table-set!
             count-table
             hash-key
             (+ (hash-table-ref count-table hash-key (λ () 0)) 1))
            (loop (+ times-tried 1)))]
         [else
          count-table])))))


(define simulate-fight
  (lambda* (fight-situation
            rand-int-proc
            #:key
            (rules *default-risk-rules*))
    "Calculate a fight result for rolling dice as long as
possible for either side, always rolling the maximum number
of dice possible."
    (let loop ([prev-round-result
                (make-init-fight-situation (get-att-rems fight-situation)
                                           (get-def-rems fight-situation))])
      (let ([att-rems (get-att-rems prev-round-result)]
            [def-rems (get-def-rems prev-round-result)])
        (cond
         ;; The attacker must have at least 1 unit remaining, so that the source
         ;; area remains occupied. The defender has no such need, as they can lose
         ;; the target area.
         ;; TODO: This could be subject to rule changes.
         [(or (= att-rems 1) (= def-rems 0))
          prev-round-result]

         [(and (> att-rems 1) (> def-rems 0))
          (loop
           (simulate-round prev-round-result
                           rand-int-proc
                           #:rules rules))]
         [else
          (error "illegal situation - your code is buggy")])))))
